from typing import List
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
# from langchain.schema import BaseOutputParser
from langchain.output_parsers import CommaSeparatedListOutputParser
import re
from PyCmpltrtok.common import sep

llm = ChatOpenAI(
    streaming=True,
    verbose=True,
    callbacks=[],
    openai_api_key="token1",
    openai_api_base=f"http://127.0.0.1:6001/v1",
    model_name="chatglm2-6b-int4",
    temperature=0.0,
    openai_proxy=None,
    top_p=1.0,
    max_tokens=2048,
)

# if 0:
#     class CommaSeparatedListOutputParser(BaseOutputParser[List[str]]):
#         """Parse the output of an LLM call to a comma-separated list."""

#         def parse(self, text: str) -> List[str]:
#             """Parse the output of an LLM call."""
#             # return text.strip().split(", ")
#             return re.split(r',\s*', text.strip())


template = """You are a helpful assistant who generates comma separated lists.
A user will pass in a category, and you should generate 5 objects in that category in a comma separated list.
ONLY return a comma separated list, and nothing more."""
template = """你是一个生成逗号分隔列表的助手。
用户将提供一个类别，你生成这个类别的5个对象组成的逗号分隔列表。
只生成逗号分隔列表，不生成其他内容。
"""
human_template = "{text}"

chat_prompt = ChatPromptTemplate.from_messages([
    ("system", template),
    ("human", human_template),
])

category = "颜色"  # colors, presidents, 颜色, 汽车, 各种汽车

sep('prompt | llm | parser')
chain = chat_prompt | llm | CommaSeparatedListOutputParser()
gen = chain.stream({"text": category})
for i, result in enumerate(gen):
    print(i, f'|{result}|')
# >> ['red', 'blue', 'green', 'yellow', 'orange']

sep('prompt | llm')
chain = chat_prompt | llm
gen = chain.stream({"text": category})
for i, result in enumerate(gen):
    print(i, f'|{result}|')